import logging
import time

import torch

from gxl_ai_utils.utils import utils_file


import utils_local
# 目录转移
# input_dir = "/home/backup_nfs2/nfs1_data/data_aishell"
# output_dir = "/home/work_nfs14/asr_data/data/aishell_1/origin_wav/data_aishell"

utils_file.do_config_root_logger()
# logging.info("开始转移origin wav dir")
# now = time.time()
# utils_file.do_copy_dir(input_dir, output_dir)
# logging.info(f"结束转移origin wav dir, 耗时{time.time() - now}秒")



logging.info('开始生成raw scp')
output_dir = "/home/work_nfs15/asr_data/data/aishell_1/raw_scp"
input_dir = "/home/work_nfs15/asr_data/data/aishell_1/origin_wav/data_aishell/transcript"
partions = ['dev', 'test', 'train']
now = time.time()
for partion in partions:
    input_dir_tmp = f"{input_dir}/{partion}"
    output_dir_tmp = f"{output_dir}/{partion}"
    text_path = f"{input_dir_tmp}/text"
    output_text_path = f"{output_dir_tmp}/text"
    text_dict = utils_file.load_dict_from_scp(text_path)
    new_text_dict = {}
    for key, value in text_dict.items():
        new_value = value.replace(' ', '')
        new_text_dict[key] = new_value
    utils_file.write_dict_to_scp(new_text_dict, output_text_path)
    wav_scp_path = f"{input_dir_tmp}/wav.scp"
    output_wav_scp_path = f"{output_dir_tmp}/wav.scp"
    wav_dict = utils_file.load_dict_from_scp(wav_scp_path)
    new_wav_dict = {}
    for key, value in wav_dict.items():
        new_value = value.replace('/home/work_nfs/common/data/data_aishell/wav/', '/home/work_nfs15/asr_data/data/aishell_1/origin_wav/data_aishell/wav/')
        new_wav_dict[key] = new_value
    utils_file.write_dict_to_scp(new_wav_dict, output_wav_scp_path)
    output_data_list_path = f"{output_dir_tmp}/data.list"
    utils_file.do_convert_wav_text_scp_to_jsonl(output_wav_scp_path, output_text_path, output_data_list_path)



logging.info(f"结束生成raw scp, 耗时{time.time() - now}秒")

# logging.info('开始打shards')
# input_dir  = "/home/work_nfs14/asr_data/data/aishell_1/raw_scp"
# output_dir = "/home/work_nfs14/asr_data/data/aishell_1/shards_file"
# partions = ['dev', 'test', 'train']
# now = time.time()
# for partion in partions:
#     input_dir_tmp = f"{input_dir}/{partion}"
#     output_dir_tmp = f"{output_dir}/{partion}"
#     utils_file.makedir_sil(output_dir_tmp)
#     wav_scp_path = f"{input_dir_tmp}/wav.scp"
#     text_scp_path = f"{input_dir_tmp}/text"
#     utils_file.do_make_shard_file(wav_scp_path, text_scp_path, output_dir_tmp)
# logging.info(f'结束打shards, 耗时{time.time() - now}秒')

# logging.info(f'开始提取k2特征')
# if __name__== "__main__":
#     torch.set_num_threads(1)
#     torch.set_num_interop_threads(1)
#     input_dir = "/home/work_nfs14/asr_data/data/aishell_1/raw_scp"
#     output_dir = "/home/work_nfs14/asr_data/data/aishell_1/k2_format"
#     partions = ['dev', 'test', 'train']
#     now = time.time()
#     for partion in partions:
#         input_dir_tmp = f"{input_dir}/{partion}"
#         output_dir_tmp = f"{output_dir}/{partion}"
#         utils_file.makedir_sil(output_dir_tmp)
#         wav_scp_path = f"{input_dir_tmp}/wav.scp"
#         text_scp_path = f"{input_dir_tmp}/text"
#         utils_file.do_make_data4icefall(wav_scp_path, text_scp_path, parent_dir=output_dir_tmp, partition=partion,
#                                         prefix='aishell')
#     logging.info(f'结束打提取特征, 耗时{time.time() - now}秒')
